from kafka import KafkaProducer
from kafka import KafkaConsumer
from kafka.errors import KafkaError
import pandas as pd


class Kafka_producer():

    """
    生产模块：根据不同的key，区分消息
    """

    def __init__(self, bootstrap_servers, kafkatopic, key=None):
        self.bootstrap_servers = bootstrap_servers
        self.kafkatopic = kafkatopic
        self.key = key
        self.producer = KafkaProducer(bootstrap_servers=self.bootstrap_servers)

    def send_df(self, df,topic=None):
        try:
            if topic is None:
                topic = self.kafkatopic
            message = df
            producer = self.producer
            producer.send(topic, key=self.key, value=message.to_msgpack())  # 注意转换
            producer.flush()
        except KafkaError as e:
            print(e)


class Kafka_consumer():

    def __init__(self,bootstrap_servers, kafkatopic, groupid=None, key=None):
        self.bootstrap_servers = bootstrap_servers
        self.kafkatopic = kafkatopic
        self.groupid = groupid
        self.key = key
        self.consumer = KafkaConsumer(self.kafkatopic, group_id=self.groupid,bootstrap_servers=self.bootstrap_servers)

    def read_df(self):
        try:
            for message in self.consumer:
                yield pd.read_msgpack(message.value) #返回迭代器
        except Exception as e:
            print(e)
